#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import tempfile
import re
import logging
import uuid
from typing import Dict, List, Any, BinaryIO, Tuple, Optional
from pathlib import Path
import shutil
import io
from datetime import datetime

from docx import Document
from PIL import Image

from document_processor.engines.base_engine import BaseEngine

logger = logging.getLogger(__name__)

class WordEngine(BaseEngine):
    """使用 python-docx 处理 Word 文档的引擎"""
    
    def __init__(self, config: Dict[str, Any]):
        """初始化 Word 处理引擎
        
        Args:
            config: 配置字典
        """
        super().__init__(config)
        logger.info("已初始化 Word 处理引擎")

    async def convert_to_markdown(
        self, 
        file_content: BinaryIO, 
        file_name: str,
        image_option: str = "obs",
        work_dir: Optional[str] = None,
        md_filename: Optional[str] = None
    ) -> Dict[str, Any]:
        """使用Word处理引擎将Word文档转换为Markdown格式
        
        Args:
            file_content: 文件内容（二进制流）
            file_name: 文件名
            image_option: 图片处理选项，可选值："obs"（上传到OBS）或"remove"（移除图片）
            work_dir: 工作目录路径，如果提供则使用此目录，否则创建临时目录
            md_filename: 生成的Markdown文件名（不含扩展名），用于命名图片目录
            
        Returns:
            包含以下键的字典:
                - text: 提取的纯文本
                - markdown: 转换后的Markdown文本
                - images: 图片信息列表
                - temp_dir: 临时目录路径
        """
        # 验证图片处理选项
        if image_option not in ["obs", "remove"]:
            logger.warning(f"不支持的图片处理选项: {image_option}，使用 obs 选项")
            image_option = "obs"
        
        # 检查文件格式
        file_ext = os.path.splitext(file_name)[1].lower()
        if file_ext not in ['.doc', '.docx']:
            raise ValueError(f"Word引擎只支持DOC/DOCX格式，不支持: {file_ext}")
            
        # 使用提供的工作目录或创建临时目录
        temp_dir = work_dir if work_dir else tempfile.mkdtemp()
        if not os.path.exists(temp_dir):
            os.makedirs(temp_dir)
            
        # 确保output目录存在
        output_dir = os.path.join(temp_dir, "output")
        if not os.path.exists(output_dir):
            os.makedirs(output_dir)
            
        temp_file_path = os.path.join(temp_dir, file_name)
        
        try:
            # 保存上传的文件到临时位置
            file_content_bytes = file_content.read()
            with open(temp_file_path, "wb") as f:
                f.write(file_content_bytes)
            
            # 使用 python-docx 读取文档
            logger.info(f"开始转换Word文档: {file_name}")
            
            doc = Document(temp_file_path)
            
            # 第1步：简单分析文档结构，构建内容列表（段落、表格、图片）
            content_elements = []
            
            # 创建图片文件夹（使用Markdown文件名作为目录名）
            if md_filename:
                images_dir_name = f"{md_filename}_images"
            else:
                # 如果没有提供Markdown文件名，则使用UUID
                dir_id = uuid.uuid4().hex[:8]
                images_dir_name = f"doc_images_{dir_id}"
                
            images_dir = os.path.join(output_dir, images_dir_name)
            os.makedirs(images_dir, exist_ok=True)
            
            # 首先提取文档中的所有图片
            images = []
            image_mappings = {}  # 图片内部ID到索引的映射
            
            if image_option != "remove":
                # 提取所有图片
                image_index = 0
                
                # 处理文档中的图片关系
                for rel_id, rel in doc.part.rels.items():
                    if "image" in rel.target_ref:
                        try:
                            # 获取图片数据
                            image_data = rel.target_part.blob
                            
                            # 检查图片尺寸
                            img = Image.open(io.BytesIO(image_data))
                            width, height = img.size
                            
                            # 忽略太小的图片（可能是页眉页脚或装饰性图片）
                            if width < 50 or height < 50:
                                continue
                                
                            # 生成唯一文件名，但保持有序性
                            img_format = img.format.lower() if img.format else "png"
                            img_filename = f"image_{image_index:03d}.{img_format}"
                            
                            # 记录图片信息
                            image_info = {
                                "filename": img_filename,
                                "format": img_format,
                                "width": width,
                                "height": height,
                                "rel_id": rel_id,
                                "dir_name": images_dir_name,  # 记录图片所在目录名
                                "data": image_data  # 保存图片二进制数据
                            }
                            
                            images.append(image_info)
                            
                            # 创建映射，便于后续查找
                            image_mappings[rel_id] = image_index
                            image_index += 1
                        except Exception as e:
                            logger.error(f"处理图片时出错: {str(e)}")
            
            # 第2步：按顺序处理所有元素，构建Markdown文档
            # 先分析每个段落中包含的图片信息
            for para in doc.paragraphs:
                para_content = {"type": "paragraph", "text": para.text, "style": para.style.name, "images": []}
                
                # 检查段落中是否包含图片
                for run in para.runs:
                    if not hasattr(run, '_element'):
                        continue
                    
                    # 查找所有drawing元素
                    try:
                        for drawing_ns in ['{http://schemas.openxmlformats.org/wordprocessingml/2006/main}drawing',
                                          '{http://schemas.openxmlformats.org/drawingml/2006/main}drawing']:
                            for drawing in run._element.findall(f'.//{drawing_ns}', namespaces=run._element.nsmap):
                                # 在drawing中查找所有blip元素（图片引用）
                                for blip_ns in ['{http://schemas.openxmlformats.org/drawingml/2006/main}blip',
                                               '{http://schemas.openxmlformats.org/drawingml/2006/picture}blip']:
                                    for blip in drawing.findall(f'.//{blip_ns}', namespaces=run._element.nsmap):
                                        # 获取图片关系ID
                                        embed_attr = '{http://schemas.openxmlformats.org/officeDocument/2006/relationships}embed'
                                        rel_id = blip.get(embed_attr)
                                        
                                        if rel_id and rel_id in image_mappings:
                                            img_idx = image_mappings[rel_id]
                                            para_content["images"].append(img_idx)
                    except Exception as e:
                        logger.warning(f"分析段落中的图片时出错: {str(e)}")
                
                # 只处理有内容的段落（文本或图片）
                if para_content["text"].strip() or para_content["images"]:
                    content_elements.append(para_content)
            
            # 处理表格
            for table in doc.tables:
                table_content = {"type": "table", "rows": []}
                
                for row in table.rows:
                    row_cells = []
                    for cell in row.cells:
                        row_cells.append(cell.text.strip())
                    table_content["rows"].append(row_cells)
                
                content_elements.append(table_content)
            
            # 第3步：将内容转换为Markdown
            markdown_lines = []
            processed_images = set()
            
            for element in content_elements:
                if element["type"] == "paragraph":
                    text = element["text"].strip()
                    style = element["style"].lower()
                    
                    # 处理标题
                    if text and 'heading' in style:
                        try:
                            level = int(style.replace('heading ', ''))
                            markdown_lines.append(f"{'#' * level} {text}\n")
                        except:
                            # 如果无法解析标题级别，作为普通段落处理
                            markdown_lines.append(f"{text}\n\n")
                    elif text:
                        # 普通段落
                        markdown_lines.append(f"{text}\n\n")
                    
                    # 处理段落中的图片
                    for img_idx in element["images"]:
                        if img_idx < len(images):
                            markdown_lines.append(f"\n__IMAGE_URL_{img_idx}__\n\n")
                            processed_images.add(img_idx)
                
                elif element["type"] == "table":
                    rows = element["rows"]
                    if not rows:
                        continue
                        
                    # 处理表头
                    header_row = rows[0]
                    markdown_lines.append('| ' + ' | '.join(header_row) + ' |')
                    markdown_lines.append('|' + '|'.join(['---' for _ in header_row]) + '|')
                    
                    # 处理数据行
                    for row in rows[1:]:
                        markdown_lines.append('| ' + ' | '.join(row) + ' |')
                    
                    markdown_lines.append('\n')
            
            # 处理未被引用的图片，放到文档最后
            unprocessed_images = [i for i in range(len(images)) if i not in processed_images]
            if unprocessed_images:
                markdown_lines.append("\n## 文档图片\n\n")
                for img_idx in unprocessed_images:
                    markdown_lines.append(f"\n__IMAGE_URL_{img_idx}__\n\n")
            
            # 合并为最终的Markdown文本
            markdown_text = ''.join(markdown_lines)
            
            # 保存Markdown到临时文件
            output_file = os.path.join(output_dir, f"{Path(file_name).stem}.md")
            with open(output_file, 'w', encoding='utf-8') as f:
                f.write(markdown_text)
            
            return {
                "text": re.sub(r'!\[.*?\]\(.*?\)', '', markdown_text),  # 纯文本（移除图片引用）
                "markdown": markdown_text,
                "images": images,
                "temp_dir": temp_dir,
                "output_md": output_file,
                "images_dir": images_dir,
                "md_filename": Path(file_name).stem  # 添加Markdown文件名（不含扩展名）
            }
            
        except Exception as e:
            logger.error(f"处理文件出错: {str(e)}")
            import traceback
            logger.error(f"详细错误: {traceback.format_exc()}")
            # 不再自动清理临时目录，让用户可以访问生成的图片资源
            # if not work_dir and temp_dir and os.path.exists(temp_dir):
            #     shutil.rmtree(temp_dir, ignore_errors=True)
            raise

    async def extract_images(
        self, 
        file_content: BinaryIO,
        file_name: str
    ) -> Tuple[List[Dict[str, Any]], Optional[str]]:
        """从Word文档中提取图片
        
        Args:
            file_content: 文件内容（二进制流）
            file_name: 文件名
            
        Returns:
            Tuple[图片信息列表, 临时目录路径]
        """
        # 创建临时目录
        temp_dir = tempfile.mkdtemp()
        temp_file_path = os.path.join(temp_dir, file_name)
        
        try:
            # 保存上传的文件到临时位置
            file_content_bytes = file_content.read()
            with open(temp_file_path, "wb") as f:
                f.write(file_content_bytes)
            
            # 使用 python-docx 读取文档
            doc = Document(temp_file_path)
            images = []
            
            # 创建图片保存目录
            img_dir = os.path.join(temp_dir, "images")
            os.makedirs(img_dir, exist_ok=True)
            
            # 提取所有图片
            image_index = 0
            for rel in doc.part.rels.values():
                if "image" in rel.target_ref:
                    try:
                        # 获取图片数据
                        image_data = rel.target_part.blob
                        
                        # 检查图片尺寸
                        img = Image.open(io.BytesIO(image_data))
                        width, height = img.size
                        
                        # 忽略太小的图片
                        if width < 50 or height < 50:
                            continue
                            
                        # 生成唯一文件名
                        img_format = img.format.lower() if img.format else "png"
                        img_filename = f"image_{image_index}_{uuid.uuid4().hex}.{img_format}"
                        img_path = os.path.join(img_dir, img_filename)
                        
                        # 保存图片
                        with open(img_path, "wb") as img_file:
                            img_file.write(image_data)
                        
                        # 记录图片信息
                        images.append({
                            "path": img_path,
                            "filename": img_filename,
                            "format": img_format,
                            "width": width,
                            "height": height,
                            "rel_id": rel.rId
                        })
                        
                        image_index += 1
                    except Exception as e:
                        logger.error(f"提取图片时出错: {str(e)}")
            
            return images, temp_dir
            
        except Exception as e:
            logger.error(f"从Word文档提取图片时出错: {str(e)}")
            if temp_dir and os.path.exists(temp_dir):
                shutil.rmtree(temp_dir)
            raise

    async def handle_images(
        self, 
        images: List[Dict[str, Any]], 
        markdown_text: str,
        uploader,
        md_filename: Optional[str] = None
    ) -> str:
        """处理图片（上传到OBS并更新Markdown中的图片引用）
        
        Args:
            images: 图片信息列表
            markdown_text: Markdown文本
            uploader: OBS上传器实例
            md_filename: Markdown文件名（不含扩展名），用于生成OBS路径
            
        Returns:
            更新后的Markdown文本（图片引用已替换为OBS URL）
        """
        if not images:
            logger.info("没有图片需要处理")
            return markdown_text
        
        # 如果没有提供OBS上传器，则无法处理图片
        if not uploader or not hasattr(uploader, 'upload'):
            logger.warning("未提供有效的OBS上传器，无法处理图片")
            return markdown_text
        
        logger.info(f"开始处理 {len(images)} 张图片")
        
        # 首先检查Markdown中是否已经包含完整的图片链接格式，如果有，可能是之前已经处理过
        if re.search(r'!\[.*?\]\([^_].*?\)', markdown_text):
            logger.warning("检测到Markdown中已包含完整的图片链接格式，将跳过处理避免重复")
            return markdown_text
        
        try:
            # 解析所有占位符
            placeholder_pattern = r'__IMAGE_URL_(\d+)__'
            all_placeholders = re.findall(placeholder_pattern, markdown_text)
            logger.info(f"在Markdown中找到 {len(all_placeholders)} 个图片占位符")
            
            # 为每张图片获取URL（上传到OBS）
            image_urls = {}
            
            # 生成OBS路径前缀
            # 总是使用当前时间戳生成路径前缀，确保图片路径唯一性
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            obs_path_prefix = ""
            if md_filename:
                # 清理文件名，移除不安全的字符
                safe_filename = re.sub(r'[^\w\-\.]', '_', md_filename)
                obs_path_prefix = f"docs/{timestamp}/{safe_filename}"
            else:
                # 即使没有文件名，也使用时间戳创建唯一路径
                obs_path_prefix = f"docs/{timestamp}/document"
            logger.info(f"使用OBS路径前缀: {obs_path_prefix}")
            
            for i, image in enumerate(images):
                try:
                    # 检查图片是否有效
                    if not image or "data" not in image:
                        logger.warning(f"跳过无效的图片: {image.get('filename', '未知')}")
                        continue
                        
                    logger.info(f"处理图片 {i+1}/{len(images)}: {image['filename']}")
                    
                    # 上传图片到OBS
                    try:
                        # 准备上传数据
                        upload_data = {
                            "filename": image['filename'],
                            "data": image['data'],
                            "content_type": f"image/{image['format']}"
                        }
                        
                        # 构建OBS路径
                        # 格式: 文档名_图片序号_原始文件名
                        doc_name = md_filename if md_filename else "doc"
                        # 清理文档名，移除不安全的字符
                        safe_doc_name = re.sub(r'[^\w\-\.]', '_', doc_name)
                        # 获取图片序号（从1开始）
                        img_number = i + 1
                        # 获取原始文件名（不含扩展名）
                        original_name = os.path.splitext(image['filename'])[0]
                        # 获取扩展名
                        ext = os.path.splitext(image['filename'])[1]
                        # 构建新文件名
                        new_filename = f"{safe_doc_name}_img{img_number:03d}{ext}"
                        
                        # 构建完整的OBS路径
                        obs_path = f"{obs_path_prefix}/{new_filename}"
                        upload_data["object_key"] = obs_path  # 使用object_key替代path
                        logger.info(f"使用OBS路径: {obs_path}")
                        
                        # 调用上传方法并等待完成
                        img_url = await uploader.upload(upload_data)
                        logger.info(f"图片上传成功: {img_url}")
                        
                        if not img_url:
                            logger.warning(f"图片 {image['filename']} 上传返回空URL")
                            continue
                        
                    except Exception as e:
                        logger.error(f"上传图片失败: {str(e)}")
                        continue
                    
                    # 使用占位符作为键
                    placeholder = f"__IMAGE_URL_{i}__"
                    image_urls[placeholder] = img_url
                    logger.info(f"处理图片: {placeholder} -> {img_url}")
                    
                except Exception as e:
                    logger.error(f"处理图片 {image.get('filename', '未知')} 时出错: {str(e)}")
                    # 继续处理下一张图片
                    continue
            
            # 一次性替换所有占位符，避免多次字符串替换
            processed_markdown = markdown_text
            
            # 创建正则表达式替换模式
            def replace_placeholder(match):
                """替换占位符为Markdown图片链接"""
                placeholder = match.group(0)
                if placeholder not in image_urls:
                    return placeholder  # 未找到对应的URL，保留原占位符
                
                url = image_urls[placeholder]
                
                # 获取图片索引
                img_idx = int(match.group(1))
                img_caption = f"图片_{img_idx+1}"
                
                # 不添加尺寸限制，避免影响图片预览
                return f"![{img_caption}]({url})"
            
            # 使用正则替换
            processed_markdown = re.sub(placeholder_pattern, replace_placeholder, processed_markdown)
            
            # 检查是否还有未替换的占位符
            remaining_placeholders = re.findall(placeholder_pattern, processed_markdown)
            if remaining_placeholders:
                logger.warning(f"还有 {len(remaining_placeholders)} 个未替换的占位符")
                
                # 只清理独立成行的未替换占位符，保留内嵌在文本中的占位符（可能有语义）
                lines = processed_markdown.split('\n')
                filtered_lines = []
                
                for line in lines:
                    # 如果一行只包含占位符，则跳过此行
                    if re.match(r'^\s*__IMAGE_URL_\d+__\s*$', line):
                        logger.debug(f"移除独立占位符行: {line.strip()}")
                        continue
                    filtered_lines.append(line)
                
                processed_markdown = '\n'.join(filtered_lines)
            else:
                logger.info("所有占位符都已替换")
            
            logger.info(f"图片处理完成，共处理 {len(images)} 张图片")
            return processed_markdown
            
        except Exception as e:
            logger.error(f"处理图片集合时出错: {str(e)}")
            import traceback
            logger.error(f"详细错误: {traceback.format_exc()}")
            return markdown_text 